"""Retry mechanisms for AI service."""

import asyncio
import logging
import random
from typing import TypeVar, Callable, Optional, Type, Tuple, Any
from functools import wraps
from datetime import datetime, timedelta

from ..exceptions import (
    RateLimitError,
    TimeoutError,
    ModelError,
    ModelNotAvailableError
)

logger = logging.getLogger(__name__)

T = TypeVar('T')


class RetryStrategy:
    """Base retry strategy class."""
    
    def __init__(
        self,
        max_retries: int = 3,
        base_delay: float = 1.0,
        max_delay: float = 60.0
    ):
        """Initialize retry strategy.
        
        Args:
            max_retries: Maximum number of retry attempts
            base_delay: Base delay in seconds
            max_delay: Maximum delay in seconds
        """
        self.max_retries = max_retries
        self.base_delay = base_delay
        self.max_delay = max_delay
    
    def get_delay(self, attempt: int) -> float:
        """Get delay for given attempt.
        
        Args:
            attempt: Attempt number (0-based)
            
        Returns:
            Delay in seconds
        """
        raise NotImplementedError


class ExponentialBackoff(RetryStrategy):
    """Exponential backoff retry strategy."""
    
    def __init__(
        self,
        max_retries: int = 3,
        base_delay: float = 1.0,
        max_delay: float = 60.0,
        multiplier: float = 2.0,
        jitter: bool = True
    ):
        """Initialize exponential backoff.
        
        Args:
            max_retries: Maximum number of retry attempts
            base_delay: Base delay in seconds
            max_delay: Maximum delay in seconds
            multiplier: Delay multiplier for each attempt
            jitter: Whether to add random jitter
        """
        super().__init__(max_retries, base_delay, max_delay)
        self.multiplier = multiplier
        self.jitter = jitter
    
    def get_delay(self, attempt: int) -> float:
        """Get delay with exponential backoff.
        
        Args:
            attempt: Attempt number (0-based)
            
        Returns:
            Delay in seconds
        """
        delay = min(
            self.base_delay * (self.multiplier ** attempt),
            self.max_delay
        )
        
        if self.jitter:
            # Add random jitter (0-25% of delay)
            delay += random.uniform(0, delay * 0.25)
        
        return delay


class LinearBackoff(RetryStrategy):
    """Linear backoff retry strategy."""
    
    def __init__(
        self,
        max_retries: int = 3,
        base_delay: float = 1.0,
        max_delay: float = 60.0,
        increment: float = 1.0
    ):
        """Initialize linear backoff.
        
        Args:
            max_retries: Maximum number of retry attempts
            base_delay: Base delay in seconds
            max_delay: Maximum delay in seconds
            increment: Delay increment for each attempt
        """
        super().__init__(max_retries, base_delay, max_delay)
        self.increment = increment
    
    def get_delay(self, attempt: int) -> float:
        """Get delay with linear backoff.
        
        Args:
            attempt: Attempt number (0-based)
            
        Returns:
            Delay in seconds
        """
        return min(
            self.base_delay + (self.increment * attempt),
            self.max_delay
        )


def retry_with_backoff(
    strategy: Optional[RetryStrategy] = None,
    retryable_exceptions: Tuple[Type[Exception], ...] = (
        RateLimitError,
        TimeoutError,
        ModelError
    ),
    fallback: Optional[Callable] = None
):
    """Decorator for retrying functions with backoff.
    
    Args:
        strategy: Retry strategy to use
        retryable_exceptions: Tuple of exceptions to retry
        fallback: Optional fallback function
        
    Returns:
        Decorated function
    """
    if strategy is None:
        strategy = ExponentialBackoff()
    
    def decorator(func: Callable[..., T]) -> Callable[..., T]:
        @wraps(func)
        async def async_wrapper(*args, **kwargs) -> T:
            last_exception = None
            
            for attempt in range(strategy.max_retries + 1):
                try:
                    return await func(*args, **kwargs)
                    
                except retryable_exceptions as e:
                    last_exception = e
                    
                    if attempt >= strategy.max_retries:
                        logger.error(
                            f"Max retries ({strategy.max_retries}) exceeded for {func.__name__}"
                        )
                        break
                    
                    # Check for rate limit with specific retry_after
                    if isinstance(e, RateLimitError) and e.retry_after:
                        delay = e.retry_after
                    else:
                        delay = strategy.get_delay(attempt)
                    
                    logger.warning(
                        f"Attempt {attempt + 1}/{strategy.max_retries + 1} failed: {e}. "
                        f"Retrying in {delay:.1f}s..."
                    )
                    
                    await asyncio.sleep(delay)
                    
                except Exception as e:
                    # Non-retryable exception
                    logger.error(f"Non-retryable error in {func.__name__}: {e}")
                    raise
            
            # All retries failed
            if fallback:
                logger.info(f"Using fallback for {func.__name__}")
                return await fallback(*args, **kwargs) if asyncio.iscoroutinefunction(fallback) else fallback(*args, **kwargs)
            
            raise last_exception
        
        @wraps(func)
        def sync_wrapper(*args, **kwargs) -> T:
            last_exception = None
            
            for attempt in range(strategy.max_retries + 1):
                try:
                    return func(*args, **kwargs)
                    
                except retryable_exceptions as e:
                    last_exception = e
                    
                    if attempt >= strategy.max_retries:
                        logger.error(
                            f"Max retries ({strategy.max_retries}) exceeded for {func.__name__}"
                        )
                        break
                    
                    if isinstance(e, RateLimitError) and e.retry_after:
                        delay = e.retry_after
                    else:
                        delay = strategy.get_delay(attempt)
                    
                    logger.warning(
                        f"Attempt {attempt + 1}/{strategy.max_retries + 1} failed: {e}. "
                        f"Retrying in {delay:.1f}s..."
                    )
                    
                    import time
                    time.sleep(delay)
                    
                except Exception as e:
                    logger.error(f"Non-retryable error in {func.__name__}: {e}")
                    raise
            
            if fallback:
                logger.info(f"Using fallback for {func.__name__}")
                return fallback(*args, **kwargs)
            
            raise last_exception
        
        return async_wrapper if asyncio.iscoroutinefunction(func) else sync_wrapper
    
    return decorator


class ModelFallback:
    """Model fallback handler for automatic degradation."""
    
    def __init__(self, models: list, primary_model: Optional[str] = None):
        """Initialize model fallback handler.
        
        Args:
            models: List of available models in priority order
            primary_model: Primary model to use
        """
        self.models = models
        self.primary_model = primary_model or models[0] if models else None
        self.current_model_index = 0
        self.failure_counts = {model: 0 for model in models}
        self.last_failure_times = {model: None for model in models}
        self.cooldown_period = timedelta(minutes=5)
        self.max_failures = 3
    
    def get_current_model(self) -> Optional[str]:
        """Get current model to use.
        
        Returns:
            Model name or None
        """
        if not self.models:
            return None
        
        # Check if primary model is available
        if self.primary_model and self._is_model_available(self.primary_model):
            return self.primary_model
        
        # Find next available model
        for model in self.models:
            if self._is_model_available(model):
                return model
        
        # If all models are in cooldown, use the one with oldest failure
        oldest_model = min(
            self.models,
            key=lambda m: self.last_failure_times.get(m) or datetime.min
        )
        return oldest_model
    
    def record_failure(self, model: str):
        """Record a model failure.
        
        Args:
            model: Name of the failed model
        """
        self.failure_counts[model] += 1
        self.last_failure_times[model] = datetime.utcnow()
        
        logger.warning(
            f"Model '{model}' failed. Failure count: {self.failure_counts[model]}"
        )
    
    def record_success(self, model: str):
        """Record a model success.
        
        Args:
            model: Name of the successful model
        """
        # Reset failure count on success
        self.failure_counts[model] = 0
        self.last_failure_times[model] = None
    
    def _is_model_available(self, model: str) -> bool:
        """Check if a model is available.
        
        Args:
            model: Model name
            
        Returns:
            True if model is available
        """
        # Check failure count
        if self.failure_counts.get(model, 0) >= self.max_failures:
            # Check if cooldown period has passed
            last_failure = self.last_failure_times.get(model)
            if last_failure:
                if datetime.utcnow() - last_failure < self.cooldown_period:
                    return False
                else:
                    # Reset after cooldown
                    self.failure_counts[model] = 0
                    self.last_failure_times[model] = None
        
        return True
    
    async def execute_with_fallback(
        self,
        func: Callable,
        *args,
        **kwargs
    ) -> Any:
        """Execute function with automatic model fallback.
        
        Args:
            func: Function to execute
            *args: Positional arguments
            **kwargs: Keyword arguments
            
        Returns:
            Function result
        """
        last_exception = None
        
        for _ in range(len(self.models)):
            model = self.get_current_model()
            if not model:
                raise ModelNotAvailableError("No models available", self.models)
            
            try:
                # Execute with current model
                kwargs["model_name"] = model
                result = await func(*args, **kwargs) if asyncio.iscoroutinefunction(func) else func(*args, **kwargs)
                
                # Record success
                self.record_success(model)
                return result
                
            except Exception as e:
                last_exception = e
                self.record_failure(model)
                logger.error(f"Model '{model}' failed: {e}")
        
        # All models failed
        raise last_exception or ModelError("All models failed")